/*
 * Copyright 2021 ZhongChao Blockchain Technology Research Institute.
 *
 * Copyright 2019 Web3 Labs Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 *
 * This file has been modified by ZhongChao Blockchain Technology Research Institute to add support for Brochain.
 */
package link.luyu.plugin.brochain.crypto;

import java.math.BigInteger;
import java.security.KeyPair;
import java.util.Arrays;
import java.util.Objects;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.web3j.utils.Numeric;

/** The SM2 generated key pair. */
public class SM2KeyPair {

    private final BigInteger privateKey;
    private final BigInteger publicKey;

    public SM2KeyPair(BigInteger privateKey, BigInteger publicKey) {
        this.privateKey = privateKey;
        this.publicKey = publicKey;
    }

    public BigInteger getPrivateKey() {
        return privateKey;
    }

    public BigInteger getPublicKey() {
        return publicKey;
    }

    /**
     * Sign a hash with the private key of this key pair.
     *
     * @param messageHash the hash to sign
     * @return An {@link SM2Signature} of the hash
     */
    public SM2Signature sign(byte[] messageHash) {
        SM2Signer signer = new SM2Signer();

        ECPrivateKeyParameters privateKeyParameters =
                new ECPrivateKeyParameters(privateKey, SM2.CURVE);
        signer.init(true, privateKeyParameters);
        return signer.generateSignature(messageHash); // do not canonicalize signature
    }

    public static SM2KeyPair create(KeyPair keyPair) {
        BCECPrivateKey privateKey = (BCECPrivateKey) keyPair.getPrivate();
        BCECPublicKey publicKey = (BCECPublicKey) keyPair.getPublic();

        BigInteger privateKeyValue = privateKey.getD();

        byte[] publicKeyBytes = publicKey.getQ().getEncoded(false);
        BigInteger publicKeyValue =
                new BigInteger(1, Arrays.copyOfRange(publicKeyBytes, 1, publicKeyBytes.length));

        return new SM2KeyPair(privateKeyValue, publicKeyValue);
    }

    public static SM2KeyPair create(BigInteger privateKey) {
        return new SM2KeyPair(privateKey, SM2.publicKeyFromPrivate(privateKey));
    }

    public static SM2KeyPair create(byte[] privateKey) {
        return create(Numeric.toBigInt(privateKey));
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }

        SM2KeyPair sm2KeyPair = (SM2KeyPair) o;

        if (!Objects.equals(privateKey, sm2KeyPair.privateKey)) {
            return false;
        }

        return Objects.equals(publicKey, sm2KeyPair.publicKey);
    }

    @Override
    public int hashCode() {
        int result = privateKey != null ? privateKey.hashCode() : 0;
        result = 31 * result + (publicKey != null ? publicKey.hashCode() : 0);
        return result;
    }

    public byte[] getAddress() {
        byte[] bts = Numeric.toBytesPadded(this.publicKey, Keys.PUBLIC_KEY_SIZE);
        byte[] hash = Hash.sm3(bts);
        return Arrays.copyOfRange(hash, hash.length - Keys.ADDRESS_SIZE / 8, hash.length);
    }
}
